"""
用于拆解用户的问题，将用户的问题拆解为若干个最小的解决方法

用户： 看一下去年平均销量大于100的产品中，今年的损耗率情况

如果我一次性编写下来，会很复杂，很困难
那么采用分而治之的思想，将复杂问题拆解为若干个简单的子问题

1、计算2024年所有产品的平均销量 为A；
2、从A中筛选出平均销量大于100的产品 为B；
3、从B中查找2025年的损耗率 为C；
"""

from nl2sql.model.llm import LLM
from typing import Optional, List
from nl2sql.prompt.decompose_query import prompt
from nl2sql.tools.extract_block import extract_python


class DecomposeQuery:
    def __init__(self,
                 model: LLM = None) -> None:

        self.prompt = prompt
        # 给这个类赋予一个属性名为llm，这个属性的数据类型是LLM，或者是None
        self.llm: Optional[LLM] = model

    def decompose_query(self, query: str) -> Optional[str]:
        """
        用于拆解用户的问题，为子问题
        :param query:
        :return: None or List of String
        """
        # 把Prompt模板，进行填充
        filled_prompt = self.prompt.format(query=query)
        # 让大模型对填充好的模板，进行推理
        llm_resp: str = self.llm.chat(prompt=filled_prompt)
        # 对大模型的输出，进行后处理（因为我们没办法保证大模型输出的格式就是我们想要的）
        result = self.postprocess(text=llm_resp)
        return result

    @staticmethod
    def postprocess(text: str) -> List[str]:
        """
        将大模型的输出，转换为规定格式的数据
        :param text:
        :return:
        """
        # 1. 抽取出用md格式包裹的代码块；
        extracted = extract_python(text=text)
        return extracted


if __name__ == '__main__':
    model_name = "qwen-max"
    api_key = "sk-34f5c792513c423f90b404d28b070f1f"
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    llm = LLM(model_name=model_name,
              api_key=api_key,
              base_url=base_url)
    dq = DecomposeQuery(llm)
    response = dq.decompose_query(query="找出地市维度中，上传速度超过10Mbps且游戏/视频/社交类流量合计超过夜间流量80%的用户中，平均下载速度排名前三的地市名称和对应平均下载速度，且这些地市的用户数量需大于所有地市用户数中位数。")
    print(response)



